import gradio as gr
import requests
import base64
from utils.utils import gen_random_job_id, logger, audio_to_base64
import os
# 创建临时目录
if not os.path.exists("tmp"):
    os.makedirs("tmp")

api_url = "http://localhost:6006"

def base64_to_audio_file(base64_string, output_path):
    audio_bytes = base64.b64decode(base64_string)
    with open(output_path, "wb") as f:
        f.write(audio_bytes)
    return output_path

class AudioDH:
    def __init__(self):
        self.system_prompt = "你是minglog，是由骆明开发的聊天机器人。你擅长编写人工智能和大模型相关的代码，有丰富的项目开发经验。"

    def asr(self, audio_path):
        api = f"{api_url}/asr"
        logger.info(f"ASR In: {audio_path}")
        resp = requests.post(api, json={"audio_path": audio_to_base64(audio_path)})
        if resp.status_code == 200 and resp.json().get("code") == 0:
            text = resp.json().get("content")
            logger.info(f"ASR Out: {text}")
            print("-" * 20)
        else:
            raise gr.Error("ASR 出错，请检查ASR服务。")
        return text

    def llm(self, text, system_message, message):
        api = f"{api_url}/llm"
        logger.info(f"LLM In: {text}")
        resp = requests.post(api, json={"text": text, "system_message": system_message, "message": message})
        if resp.status_code == 200 and resp.json().get("code") == 0:
            resp_text = resp.json().get("content")
            logger.info(f"LLM Out: {resp_text}")
            print("-" * 20)
        else:
            raise gr.Error("LLM 出错，请检查LLM服务。")
        return resp_text

    def tts(self, text):
        api = f"{api_url}/tts"
        logger.info(f"TTS In: {text}")
        resp = requests.post(api, json={"text": text})
        if resp.status_code == 200 and resp.json().get("code") == 0:
            audio_base64 = resp.json().get("content")
            save_path = f"tmp/{gen_random_job_id()}.mp3"
            base64_to_audio_file(audio_base64, save_path)
            logger.info(f"TTS Out: {save_path}")
            print("-" * 20)
            return save_path
        else:
            raise gr.Error("TTS 出错，请检查TTS服务。")
        
    # 模拟处理音频并回应的函数
    def respond(self, audio):
        text = self.asr(audio)
        llm_text = self.llm(text, self.system_prompt, [])[0]
        return_audio = self.tts(llm_text)
        response1 = gr.Audio(value=audio, autoplay=False)
        response2 = gr.Audio(value=return_audio, autoplay=True)
        return [(response1, response2)], None  # 清空输入框
    
    def create_ui(self):
        with gr.Blocks() as demo:
            chatbot = gr.Chatbot()
            with gr.Row():
                audio_input = gr.Audio(sources="microphone", type="filepath", label="上传音频")
            audio_input.stop_recording(fn=self.respond, inputs=[audio_input], outputs=[chatbot, audio_input])
            # audio_input.upload(fn=self.respond, inputs=[audio_input], outputs=[chatbot, audio_input])
        return demo
    
audio_dh = AudioDH()
demo = audio_dh.create_ui()
demo.launch(allowed_paths=["tmp"], server_name="0.0.0.0", server_port=7860)
